# calibrate.py
# Austin Abrams, Washington University in St. Louis
# 5/10/2010
    
# Given a scene, findCameraLocation(scene) will find the list of correspondences in that
# scene and infer the camera's location, orientation, and each correspondence's projected
# pixel error.

from models import *
from common import helperFunctions, matfunc, svd, constants, commonDB, Simplex
import math

# normalizes the lists into (x,y,z) space, whose origin is the first point in the lists.
def normalizeLatLon(latList, lonList, altList):
    originLat = latList[0];
    originLon = lonList[0];
    originAlt = altList[0];

    x = []
    y = []
    z = []
    
    for i in range(len(lonList)):
        x.append(helperFunctions.distance(originLat,originLon,originLat,lonList[i]))
        y.append(helperFunctions.distance(originLat,originLon,latList[i],originLon))
        
        if originLat > latList[i]:
            y[i] = -1 * y[i];

        if originLon > lonList[i]:
            x[i] = -1 * x[i];
            
        z.append(altList[i] - originAlt)

    return x, y, z
    
# calibrates the camera, given a set of (x,y,z)->(u,v) correspondences.
# described in "Introductory Techniques for 3-D Computer Vision", by
# Emanuele Trucco and Alessandro Verri.  Returns the projection M matrix.
def calibrate(x,y,z,u,v, width, height):
    # find the projection matrix M.
    A = []

    for i in range(len(x)):
        A.append([x[i], y[i], z[i], 1, 0, 0, 0, 0,          -x[i] * u[i], -y[i] * u[i], -z[i] * u[i], -u[i]]);
        A.append([0, 0, 0, 0,          x[i], y[i], z[i], 1, -x[i] * v[i], -y[i] * v[i], -z[i] * v[i], -v[i]]);

    U,D,Vt = svd.svd(A)

    # pick out the column cooresponding to the smallest value in D.
    smallestI = 0
    for i in range(len(D)):
        if D[i] < D[smallestI]:
            smallestI = i
    m = []
    for row in Vt:
        m.append(row[smallestI])

    M = [[m[0], m[1], m[2], m[3]], [m[4], m[5], m[6], m[7]], [m[8], m[9], m[10], m[11]]]

    [K,R,T] = findKRT(M)
    k33 = K[2][2]
    for i in range(len(K)):
        for j in range(len(K[i])):
            K[i][j] /= k33;
    focal = (abs(K[0][0]) + abs(K[1][1]))/2.0
    
    try:
        angles = rodrigues_matrix_to_angles(R)
    except ValueError:
        angles = [0,0,1]
        

    def params_to_matrix(solution):
        f = solution[0]
        r = solution[1:4]
        t = solution[4:7]
        K = matfunc.Matrix([[f, 0, width/2],
                            [0, -f, height/2],
                            [0, 0, 1]])
        R = matfunc.Matrix(rodrigues_angles_to_matrix(r))
        T = matfunc.Matrix([[t[0]], [t[1]], [t[2]]])
        M = K.mmul(R.augment(T))
        return [[M[0][0], M[0][1], M[0][2], M[0][3]],
                [M[1][0], M[1][1], M[1][2], M[1][3]],
                [M[2][0], M[2][1], M[2][2], M[2][3]]]
    
    # that was the least squares estimate for M, and will be a good first estimate.
    # Now, calibrate with respect to the non-linear reprojection error.
    def bundleError(solution):
        Mmat = params_to_matrix(solution)
        
        errors = [projectionError(x[i],y[i],z[i],u[i],v[i],Mmat)**2 for i in range(len(x))]
        return sum(errors)
        
    f_step = 10.0;
    r_step = 0.1;
    t_step = 10.0;
    initialGuess = [focal, angles[0], angles[1], angles[2], T[0][0], T[1][0], T[2][0]]
    increments = [f_step, r_step, r_step, r_step, t_step, t_step, t_step]
        
    s = Simplex.Simplex(bundleError, initialGuess, increments)
    values, err, iter = s.minimize(monitor=0, maxiters=10000)
    Mbundle = params_to_matrix(values)
    return Mbundle
    
def toZXZ(a,b,g):
    ca = math.cos(a)
    sa = math.sin(a)
    cb = math.cos(b)
    sb = math.sin(b)
    cg = math.cos(g)
    sg = math.sin(g)
    
    return matfunc.Matrix([[(ca * cg - cb * sa * sg), (cg * sa + ca * cb * sg), (sb * sg)],
                          [(-cb * cg * sa - ca * sg), (ca * cb * cg - sa * sg), (cg * sb)],
                          [(sa * sb), (-ca * sb), cb]])

    

# takes in a 3x3 rotation matrix and returns a possible rotation
# scheme, in (roll, tilt, heading).  This should result in the
# heading, tilt, and roll that gives points the camera toward
# the upper-left corner of the image.
def getApproxHeadingTiltRoll(R):
    r = R
    toDegrees = 180 / constants.PI
    firstRot = math.atan2(r[2][0], -r[2][1])
    secondRot = math.acos(r[2][2])
    thirdRot = math.atan2(r[0][2], r[1][2])

    firstRot *= toDegrees
    secondRot *= toDegrees
    thirdRot *= toDegrees
    
    heading = (- firstRot) % 360
    tilt    = (secondRot) % 360
    roll    = (thirdRot) % 360

    return heading, tilt, roll
    
# gets the heading and tilt that points in the direction of the
# center pixel.
def getHeadingTilt(M, width, height, heading, tilt, cameraPos):
    u = float(width) / 2.0;
    v = float(height) / 2.0;
    d = 10000.0;
    A = matfunc.Matrix([[M[0][0], M[0][1], M[0][2]],
                     [M[1][0], M[1][1], M[1][2]],
                     [M[2][0], M[2][1], M[2][2]]]);
    b = matfunc.Vec([u * d - M[0][3], v * d - M[1][3], d - M[2][3]]);
    v = A.solve(b);
    v = [v[0] - cameraPos[0], v[1] - cameraPos[1], v[2] - cameraPos[2]];
    
    toRadians = constants.PI / 180.0;
    
    # make sure this vector is pointing in the right direction.
    camTilt = tilt * toRadians;
    camHeading = heading * toRadians;
    camTheta = constants.PI - camTilt;
    camPhi = -camHeading + constants.PI / 2;
    directionX = math.sin(camTheta) * math.cos(camPhi);
    directionY = math.sin(camTheta) * math.sin(camPhi);
    directionZ = math.cos(camTheta);
    if (v[0] * directionX + v[1] * directionY + v[2] * directionZ < 0):
        v[0] = v[0] * -1;
        v[1] = v[1] * -1;
        v[2] = v[2] * -1;
    
    radius = math.sqrt(v[0]**2 + v[1]**2 + v[2]**2)
    tiltRad = math.acos(v[2]/radius)
    headingRad = math.atan2(v[1], v[0])
    heading = 90 - headingRad / toRadians;
    tilt = 180 - tiltRad / toRadians;
    return heading, tilt
    
    
def projectionError(x,y,z,u,v,M,corr=None):
    # recall that
    # [uw]   [ M00 M01 M02 M03 ]   [ X ]
    # [vw] = [ M10 M11 M12 M13 ] * [ Y ]
    # [w ]   [ M20 M21 M22 M23 ]   [ Z ]
    #                              [ 1 ]
    # So, plug in (X,Y,Z) to get the projected (u,v) coordinates, and
    # see how far off the pixel projection is.

    uw = M[0][0] * x + M[0][1] * y + M[0][2] * z + M[0][3]
    vw = M[1][0] * x + M[1][1] * y + M[1][2] * z + M[1][3]
    w  = M[2][0] * x + M[2][1] * y + M[2][2] * z + M[2][3]
    uProject = uw / w;
    vProject = vw / w;
    error = math.sqrt((uProject - u)*(uProject - u) + (vProject - v)*(vProject - v))
    
    
    if corr:
        corr.x = x;
        corr.y = y;
        corr.z = z;
        corr.pixelProjectionError = error;
        corr.reprojectedU = uProject;
        corr.reprojectedV = vProject;
        corr.put();
    return error
    
def projectionErrorL1(x,y,z,u,v,M,corr=None):
    # recall that
    # [uw]   [ M00 M01 M02 M03 ]   [ X ]
    # [vw] = [ M10 M11 M12 M13 ] * [ Y ]
    # [w ]   [ M20 M21 M22 M23 ]   [ Z ]
    #                              [ 1 ]
    # So, plug in (X,Y,Z) to get the projected (u,v) coordinates, and
    # see how far off the pixel projection is.

    uw = M[0][0] * x + M[0][1] * y + M[0][2] * z + M[0][3]
    vw = M[1][0] * x + M[1][1] * y + M[1][2] * z + M[1][3]
    w  = M[2][0] * x + M[2][1] * y + M[2][2] * z + M[2][3]
    uProject = uw / w;
    vProject = vw / w;
    error = abs(uProject - u) + abs(vProject - v);


    if corr:
        corr.x = x;
        corr.y = y;
        corr.z = z;
        corr.pixelProjectionError = error;
        corr.reprojectedU = uProject;
        corr.reprojectedV = vProject;
        corr.put();
    return error
    
# finds the projected pixel error for each correspondence in idList.
# Once the scene has been calibrated, an (X,Y,Z) point in space will
# project onto a (u,v) coordinate.  The pixel error is the distance
# betwen that (u,v) point and the 'true' (u,v) point.
def findPixelProjectionErrors(xList,yList,zList,uList,vList,M, corrs):
    
    for i in range(len(xList)):
        projectionError(xList[i], yList[i], zList[i], uList[i], vList[i], M, corrs[i])

# get the latitudes, longitudes, altitude, and image coordinates from the
# database.
def getCorrespondenceInfo(corrs):
    latList = []
    lonList = []
    altList = []
    uList = []
    vList = []
    idList = []
    for corr in corrs:
        latList.append(corr.latitude);
        lonList.append(corr.longitude);
        altList.append(corr.altitude);
        uList.append(corr.u);
        vList.append(corr.v);
        idList.append(corr.index);
    return latList, lonList, altList, uList, vList, idList


def tooFewCorrespondences(scene,corrs):
    for corr in corrs:
        if corr.pixelProjectionError:
            corr.pixelProjectionError = None
            corr.reprojectedU = None
            corr.reprojectedV = None
            corr.put()
    if scene.calibrated:
        scene.calibrated = False;
        scene.put()
    
    # see if there's a calibration that already exists for this scene.
    calib = Calibration.all().filter("scene = ", scene).get()
    if calib:
        calib.delete()

# Given the 3x4 projection matrix M, find K, the intrinsic matrix,
# R, the rotation matrix, and T, the translation vector.
def findKRT(M):
    # perform a QR decomposition to estimate the intrinsic matrix K and the
    # rotation matrix T.
    KR = [[M[0][0], M[0][1], M[0][2]], [M[1][0], M[1][1], M[1][2]], [M[2][0], M[2][1], M[2][2]]]
    KR = matfunc.Square(KR)
    [Q, U] = KR.inverse().qr()

    K = U.inverse()
    R = Q.inverse()
    
    # this QR decomp leaves some rows and columns with flipped signs.
    # change them back.

    # if R is determininant -1, flip the signs on K and R.
    for row in K:
        row[0] = -1 * row[0]
        #row[1] = -1 * row[1]

    for i in range(3):
        R[0][i] = -1 * R[0][i]
        #R[1][i] = -1 * R[1][i]


    T = K.inverse().mmul(matfunc.Matrix([[M[0][3]], [M[1][3]], [M[2][3]]]))
    return K, R, T


def calibrateCamera(scene):
    # see if we have enough correspondences.
    corrs = Correspondence.all().filter("scene = ", scene).fetch(limit=1000)
    if (len(corrs) < constants.CORRESPONDENCES_REQUIRED_FOR_CALIBRATION):
        tooFewCorrespondences(scene,corrs)
        return
    
    # get info from the correspondence list.
    latList, lonList, altList, uList, vList, idList = getCorrespondenceInfo(corrs)
    xList, yList, zList = normalizeLatLon(latList, lonList, altList);
    
    # calibrate the camera.
    M = calibrate(xList,yList,zList,uList,vList, scene.width, scene.height)
    K,R,T = findKRT(M)
    
    # estimate the position of the camera.
    pos = -1 * R.tr().mmul(T)
    cameraX = pos[0][0];
    cameraY = pos[1][0];
    cameraZ = pos[2][0];
    
    longitudePerMeter = 0.00001 / helperFunctions.distance(latList[0],lonList[0],latList[0],lonList[0] + 0.00001);
    latitudePerMeter = 0.00001 / helperFunctions.distance(latList[0],lonList[0],latList[0] + 0.00001 ,lonList[0]);
    posLat = latitudePerMeter * pos[1][0] + latList[0]
    posLon = longitudePerMeter * pos[0][0] + lonList[0]
    posAlt = pos[2][0] + altList[0]
    
    # get the heading, tilt, and roll.
    heading, tilt, roll = getApproxHeadingTiltRoll(R)
    heading, tilt = getHeadingTilt(M, scene.width, scene.height, heading, tilt, [cameraX, cameraY, cameraZ])
    if abs(roll) > 10:
        roll = 0.0
    # get the focal length.
    Kintrinsic = K / K[2][2]
    focalX = abs(Kintrinsic[0][0])
    focalY = abs(Kintrinsic[1][1])
    skew = Kintrinsic[0][1]
    imageCenterX = Kintrinsic[0][2]
    imageCenterY = Kintrinsic[1][2]
    intrinsicDivider =  K[2][2]
    
    # find pixel projection error.
    findPixelProjectionErrors(xList,yList,zList,uList,vList,M, corrs)

    # add the calibration object.
    commonDB.addOrModifyCalibration(scene, latList[0], lonList[0], altList[0], \
                           cameraX, cameraY, cameraZ, focalX, focalY, \
                           skew, imageCenterX, imageCenterY, intrinsicDivider, \
                           M, R, T)

    scene.calibrated = True
    scene.cameraLatitude = posLat
    scene.cameraLongitude = posLon
    scene.cameraAltitude = posAlt
    scene.cameraRoll = roll
    scene.cameraTilt = tilt
    scene.cameraHeading = heading
        
    scene.put()
    
    
def rodrigues_matrix_to_angles(R):
    trc = R.trace()
    trc2=(trc-1)/2;
    sinacostrc2=math.sqrt(1- trc2*trc2);
    s = [R[2][1] - R[1][2], R[0][2]-R[2][0], R[1][0]-R[0][1]];
    if (1- trc2*trc2) >= 0.000001: 
        tHeta = (math.acos(trc2));
        tHetaf=tHeta/ (2 *math.sin(tHeta));
        dtHetaf_dt=(2*math.sin(tHeta)-tHeta*2*math.cos(tHeta)) / (2 *math.sin(tHeta))/(2 *math.sin(tHeta));
    else:
        raise ValueError('input matrix to rodrigues_matrix_to_angles has determinant != -1')

    return [tHetaf * s[0], tHetaf * s[1], tHetaf * s[2]];
    
    
def rodrigues_angles_to_matrix(v):
    theta = math.sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2])
    theta2 = theta * theta;
    onecosthetatheta2=(1-math.cos(theta))/theta2; 
    sinthetatheta=math.sin(theta)/theta;
    omegav = matfunc.Matrix([[0, -v[2], v[1]], 
                             [v[2], 0, -v[0]],
                             [-v[1], v[0], 0]])
    I = matfunc.Matrix([[1., 0., 0.], 
                        [0., 1., 0.],
                        [0., 0., 1.]])
    R = I + omegav * sinthetatheta + omegav.mmul(omegav) * onecosthetatheta2
    return R